# -*- coding: utf-8 -*-
# Created Time : 2022/3/18 16:30
# Author:Zhou
# 对训练后的模型进行性能评价
import torch
from torch.utils.data import DataLoader
import wav_loader as loader
import train_utils
from si_snr import SiSnr
import matplotlib.pyplot as plt
from ANC_tools import ANR

dns_home = r"F:\NoiseData\DCCRN_data\NPP5"  # dir of dns-datas
load_batch = 4  # load batch_size(not calculate)
device = torch.device("cuda:0")  # device
batch_size = 20
criterion = SiSnr()
# 读取训练好的模型
model = torch.load('./logs/NPP5_models/parameter_epoch19_2022-03-20 16-57-12.pth')
# 打印模型
# print(model)

# 读取测试数据
train_test = train_utils.get_train_test_name(dns_home)
train_noisy_names, train_clean_names, test_noisy_names, test_clean_names = \
    train_utils.get_all_names(train_test, dns_home=dns_home)
test_dataset = loader.WavDataset(test_noisy_names, test_clean_names, frame_dur=37.5)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

for ind, (x, y) in enumerate(test_dataloader):
    x = x.view(x.size(0) * x.size(1), x.size(2)).to(device).float()
    y = y.view(y.size(0) * y.size(1), y.size(2)).to(device).float()
    for index in range(0, x.size(0) - batch_size, batch_size):
        x_item = x[index:index + batch_size, :].squeeze(0)
        y_item = y[index:index + batch_size, :].squeeze(0)
        y_p = model(x_item, train=False)
        loss = criterion(source=y_item.unsqueeze(1), estimate_source=y_p)

        x_data_cpu = x_item.data.cpu().numpy()[1, :]
        y_data_cpu = y_item.data.cpu().numpy()[1, :]
        predict_data_cpu = y_p.data.cpu().numpy().squeeze(1)[1, :]

        plt.plot(x_data_cpu)
        plt.plot(y_data_cpu)
        plt.plot(predict_data_cpu)
        plt.legend(['input', 'target', 'predict'])

        plt.show()
        print(loss)
